#ifndef __PNN_TENSOR__H
#define __PNN_TENSOR__H

#include <stdint.h>
#include <vector>
#include "define.h"

class Tensor {
   public:
    Tensor(){};
    Tensor(void *data) : m_data(data){};
    Tensor(void *data, int w = 1, int h = 1, int c = 1, int n = 1) : m_data(data), m_dims({w, h, c, n}){};
    Tensor(void *data, std::vector<int> dims) : m_data(data), m_dims(dims){};
    ~Tensor(){};

    void SetChannel(int c) { m_dims[2] = c; }

    int Width() { return m_dims.size() >= 1 ? m_dims[0] : 1; }
    int Hight() { return m_dims.size() >= 2 ? m_dims[1] : 1; }
    int Channel() { return m_dims.size() >= 3 ? m_dims[2] : 1; }
    // int Depth() { return m_dims.size() >= 3 ? m_dims[2] : 1; }
    int Batch() { return m_dims.size() >= 4 ? m_dims[3] : 1; }

    void SetDims(const std::vector<int> dims) { m_dims = dims; }
    const std::vector<int> GetDims() { return m_dims; }

    void SetData(void *data) { m_data = data; }
    void *GetData() { return m_data; }

   private:
    std::vector<int> m_dims;  // (0 : W), (1 : H), (2 : C), (3 : N)
    void *m_data = nullptr;
};

#endif
